from mcp import ClientSession, StdioServerParameters, types
import asyncio
import os

# AnyUrl用于资源URL的类型校验
from pydantic import AnyUrl

from mcp.client.stdio import stdio_client


async def on_logging(params: types.LoggingMessageNotificationParams) -> None:
    level = getattr(params, "level", "info")
    data = getattr(params, "data", None)
    logger = getattr(params, "logger", None)
    print(
        f"[LOG][{level}]",
        f"[logger][{logger}]",
        f"[data][{str(data) if data is not None else '<no-data>'}]",
    )


async def on_message(message) -> None:
    method = getattr(message, "method", None)
    if method == "notifications/progress":
        params = getattr(message, "params", None)
        if params is not None:
            progress = getattr(params, "progress", None)
            total = getattr(params, "total", None)
            message = getattr(params, "message", None)
            print(f"[PROGRESS] {progress}/{total}  - {message}")


async def main() -> None:
    base_dir = os.path.dirname(os.path.abspath(__file__))
    server_path = os.path.join(base_dir, "server.py")
    server_params = StdioServerParameters(command="python", args=[server_path], env={})
    async with stdio_client(server_params) as (read, write):
        async with ClientSession(
            read, write, logging_callback=on_logging, message_handler=on_message
        ) as session:
            await session.initialize()
            result = await session.call_tool(
                "long_running_task", {"task_name": "长时间任务", "steps": 3}
            )
            texts = []
            for block in result.content:
                if isinstance(block, types.TextContent):
                    texts.append(block.text)
            print("[long_running_task]", "|".join(texts))


if __name__ == "__main__":
    asyncio.run(main())
